#include "core/task.h"
#include "core/memory.h"
#include "core/syscall.h"
#include "cpu/cpu.h"
#include "cpu/irq.h"
#include "cpu/mmu.h" 
#include "comm/cpu_instr.h"
#include "comm/elf.h"
#include "tools/klib.h"
#include "tools/log.h"
#include "os_cfg.h"
#include "fs/fs.h"

//任务管理器
static task_manager_t task_manager;
//空闲进程的栈
static uint32_t idle_task_stack[IDLE_TASK_STACK_SIZE]; 

static task_t task_table[TASK_NR];      // 用户进程表
static mutex_t task_table_mutex;        // 进程表互斥访问锁

static int tss_init(task_t *task,int flag,uint32_t entry,uint32_t esp){
    int gdt_index=gdt_alloc_desc();
    if(gdt_index<0){
        log_printf("alloc tss failed.\n");
        return -1;
    }
    //tss初始化
    
    segment_desc_set(gdt_index, (uint32_t)&task->tss, sizeof(tss_t),
            SEG_P_PRESENT | SEG_DPL0 | SEG_TYPE_TSS);

    // tss段初始化
    kernel_memset(&task->tss, 0, sizeof(tss_t));
    
    // 需要重新分配一个特权级为0的栈，用于中断处理
    uint32_t kernel_stack=memory_alloc_page();
    if(kernel_stack ==0){
        goto tss_init_failed;    
    }

    // 根据不同的权限选择不同的访问选择子
    int code_sel, data_sel;
    if (flag & TASK_FLAG_SYSTEM) {
        code_sel = KERNEL_SELECTOR_CS;
        data_sel = KERNEL_SELECTOR_DS;
    } else {
        // 注意加了RP3,不然将产生段保护错误
        code_sel = task_manager.app_code_sel | SEG_RPL3;
        data_sel = task_manager.app_data_sel | SEG_RPL3;
    }


    task->tss.eip = entry;
    task->tss.esp =  esp;
    // 底地址到高地址
    task->tss.esp0 =kernel_stack + MEM_PAGE_SIZE;
    task->tss.ss = data_sel;
    task->tss.ss0 = KERNEL_SELECTOR_DS;
    task->tss.eip = entry;
    task->tss.eflags = EFLAGS_DEFAULT | EFLAGS_IF;
    task->tss.es =task->tss.ds= task->tss.fs = task->tss.gs = data_sel;  
    task->tss.cs = code_sel;  
    task->tss.iomap = 0;

    task->tss_sel = gdt_index;


    //为进程创建页表
    uint32_t page_table_addr= memory_create_uvm();
    if(page_table_addr <=0){
        
        goto tss_init_failed;    
    }
    task->tss.cr3=page_table_addr;

    return 0;
tss_init_failed:
    gdt_free_desc(gdt_index);
    if(kernel_stack){
        memory_free_page(kernel_stack);
    }
    return -1;
}

int task_init(task_t * task, const char * task_name,int flag,uint32_t entry,uint32_t esp){
    //对tss初始化
    ASSERT(task != (task_t *)0);
    
    //关中断
    irq_state_t state = irq_enter_protection();
    
    int err=tss_init(task,flag, entry, esp);
    if (err < 0) {
        log_printf("init task failed.\n");
        return err;
    }
    //task的属性
    kernel_strncpy(task->name,task_name,TASK_NAME_SIZE);//name
    task->slice_ticks=TASK_TIME_SLICE_DEFAULT;
    task->slice_ticks_default=TASK_TIME_SLICE_DEFAULT;
    task->sleep_ticks=0;

    //直接把task的地址作pid
    task->pid=(uint32_t)task;

    //添加到任务管理器里
    task->state=TASK_READY;

    task->heap_end=task->heap_start=0;

    //清空文件管理项
    kernel_memset(task->file_table,0,sizeof(task->file_table));

    irq_leave_protection(state);


    //log_printf("task %s node addr 0x%x",task->name,&task->list_node);
    //list_insert_last(&task_manager.task_list,&task->list_node);
    return 0;
};

void task_start(task_t * task){
    //要将task_init和插入到任务队列分开，因为sys_fork() sys_execve还会对task的状态进行操作
    irq_state_t state = irq_enter_protection();
    task_set_ready(task);
    irq_leave_protection(state);
}

void task_switch_from_to (task_t * from, task_t * to){
    //切换进程，从from  --> to
    switch_to_tss(to->tss_sel);

}

static void idle_function(void){
    //当所有进程都sleep的时候，调用这个程序
    for(;;){
        hlt();
    }
}

void task_manager_init (void){
    //任务管理器初始化
    //为运行程序的数据段，代码段分配选择子

    int sel = gdt_alloc_desc();
    segment_desc_set(sel, 0x00000000, 0xFFFFFFFF,
                     SEG_P_PRESENT | SEG_DPL3 | SEG_S_NORMAL |
                     SEG_TYPE_DATA | SEG_TYPE_RW | SEG_D);
    task_manager.app_data_sel = sel;

    sel = gdt_alloc_desc();
    segment_desc_set(sel, 0x00000000, 0xFFFFFFFF,
                     SEG_P_PRESENT | SEG_DPL3 | SEG_S_NORMAL |
                     SEG_TYPE_CODE | SEG_TYPE_RW | SEG_D);
    task_manager.app_code_sel = sel;

    //初始化分配任务表的锁
    mutex_init(&task_table_mutex);

    list_init(&task_manager.ready_list);
    list_init(&task_manager.task_list);
    list_init(&task_manager.sleep_list);
    task_manager.curr_task=(task_t *)0;
    task_init(&task_manager.idle_task,"idle task",TASK_FLAG_SYSTEM,(uint32_t)idle_function,(uint32_t)(idle_task_stack+IDLE_TASK_STACK_SIZE));
    task_start(&task_manager.idle_task);
};

void kernel_task_init(){
    void another_task_entry(void);
    uint32_t another_task_init_addr=(uint32_t) another_task_entry;
    // 以下获得的是bin文件在内存中的物理地址
    extern uint8_t s_another_task[], e_another_task[];

    uint32_t copy_size=e_another_task-s_another_task;
    
    uint32_t alloc_size = 10 * MEM_PAGE_SIZE;
    ASSERT(copy_size < alloc_size);
    //内核任务初始化
    task_init(&task_manager.kernel_task,"kernel task",0,another_task_init_addr,(another_task_init_addr+alloc_size));
    task_start(&task_manager.kernel_task);

    task_manager.kernel_task.heap_start = (uint32_t)e_another_task;  // 这里不对
    task_manager.kernel_task.heap_end = task_manager.kernel_task.heap_start;

    task_manager.curr_task=&task_manager.kernel_task;
    
    // 更新页表
    mmu_set_page_dir(task_manager.curr_task->tss.cr3);
    //为进程分配页表
    memory_alloc_page_for(another_task_init_addr,  alloc_size, PTE_P | PTE_W | PTE_U);
    //抄数据段，代码段的数据过去
    kernel_memcpy((void *)another_task_init_addr,s_another_task,copy_size);

    //把当前tr寄存器设置
    write_tr(task_manager.kernel_task.tss_sel);

};
task_t * get_kernel_task (void){
    //返回内核任务
    return &task_manager.kernel_task;
};


void task_set_ready(task_t *task){
    if(task == &task_manager.idle_task){
        return;
    }

    //把task添加到ready_list的末尾
    list_insert_last(&task_manager.ready_list,&task->list_node);
    task->state=TASK_READY;
};
void task_set_block (task_t *task){
    if(task == &task_manager.idle_task){
        return;
    }
    //把task从ready_list里删除
    list_remove(&task_manager.ready_list,&task->list_node);
    task->state=TASK_BLOCK;
};

task_t * get_task_current (void){
    //返回现在正在执行的任务
    return task_manager.curr_task;
};
void set_task_current(task_t * task){
    //设置正在执行的任务
    task_manager.curr_task=task;
};
int sys_yield (void){

    irq_state_t state = irq_enter_protection();
    //操作任务管理器
    if(list_count(&task_manager.ready_list)>1){
        task_t *curr_task=get_task_current();
        task_set_block(curr_task);
        task_set_ready(curr_task);

        task_dispatch();
    }   
    irq_leave_protection(state);
    //只有一个任务就不用切换了
    return 0;
};

task_t *get_next_ready_task(void){
    //获取下一个任务，当任务列表为0的时候，返回idle_task
    if(list_count(&task_manager.ready_list)){
        return list_node_parent(list_first(&task_manager.ready_list),task_t,list_node);
    }else{
        return &task_manager.idle_task;
    }

};

void task_dispatch (void){
    //切换任务
    task_t *curr_task=get_task_current();
    task_t *to_task=get_next_ready_task();
    if(curr_task != to_task){
        set_task_current(to_task);
        task_switch_from_to(curr_task,to_task);
    }

};

void task_time_tick(void){
   
    task_t *current_task=get_task_current();

    irq_state_t state = irq_enter_protection();
    //对正在运行的任务时间片减一，当时间片为0的时候，切换任务

    if(current_task == &task_manager.idle_task ){
        task_dispatch();
    }
    else if( --current_task->slice_ticks==0){
        //不为默认的执行函数
        current_task->slice_ticks=current_task->slice_ticks_default;
        task_set_block(current_task);
        task_set_ready(current_task);
        current_task->state=TASK_BLOCK;
        //切换任务
        task_dispatch();


    }

    if(list_count(&task_manager.sleep_list)){
            //对所有在sleep里的任务睡眠时间片减一，当时间片为0时，set_task_ready
        list_node_t * current_sleep_task_node=list_first(&task_manager.sleep_list);
        while (current_sleep_task_node)
        {
            task_t *current_task=list_node_parent(current_sleep_task_node,task_t,list_node);
            list_node_t *next_sleep_task_node=current_sleep_task_node->next;
            if(--current_task->sleep_ticks==0){
                task_set_wakeup(current_task);
            }
            current_sleep_task_node=next_sleep_task_node;
        }
    }

    

    irq_leave_protection(state);
};

void sys_sleep (uint32_t ms){
    //系统睡眠函数
    irq_state_t state=irq_enter_protection();
    task_t *current_task=get_task_current();
    task_set_sleep(current_task,ms);    
    irq_leave_protection(state);
};

uint32_t sys_getpid(void){
    return get_task_current()->pid;
};

void task_set_sleep(task_t *task, uint32_t ticks){
    //设置任务睡眠，从readylist到sleeplist
    if(ticks==0){
        return;
    }else{
        task->sleep_ticks=(ticks+(OS_TICK_MS-1))/OS_TICK_MS;
        list_remove(&task_manager.ready_list,&task->list_node);
        list_insert_last(&task_manager.sleep_list,&task->list_node);
        task->state=TASK_SLEEP;
        task_dispatch();
    }
};

void task_set_wakeup (task_t *task){
    //设置任务醒来，从sleeplist到readylist末尾
    list_remove(&task_manager.sleep_list,&task->list_node);
    list_insert_last(&task_manager.ready_list,&task->list_node);
    task->state=TASK_READY;
};


static task_t * task_alloc(){
    // 在空闲进程表格里分配一个进程
    mutex_lock(&task_table_mutex);

    task_t *target=(task_t*)0;
    for(int i=0;i<sizeof(task_table)/sizeof(task_t);i++){
        if(task_table[i].name[0]=='\0'){
            target=&task_table[i];
            break;
        }    
    }

    mutex_unlock(&task_table_mutex);

    return target;
}

static void *task_free(task_t * target){
    // 在空闲进程表格里清空一个进程
    mutex_lock(&task_table_mutex);
    target->name[0]='\0';

    mutex_unlock(&task_table_mutex);


}

/**
 * @brief 任务任务初始时,清除已经分配的资源
 */
static void task_uninit (task_t * task) {
    if (task->tss_sel) {
        // 清除对应的gdt项
        gdt_free_desc(task->tss_sel);
    }

    if (task->tss.esp0) {
        memory_free_page(task->tss.esp0 - MEM_PAGE_SIZE);
    }

    if (task->tss.cr3) {
        memory_destroy_uvm(task->tss.cr3);
    }

    kernel_memset(task, 0, sizeof(task_t));
}

uint32_t sys_fork(){
    // 根据栈，获取父进程的寄存器状态
    task_t *parent_task=get_task_current();
    task_t *child_task=task_alloc();
    if(!child_task){
        log_printf("alloc task fails\n");
        return -1;
    }


    syscall_frame_t * frame = (syscall_frame_t *)(parent_task->tss.esp0 - sizeof(syscall_frame_t));
    int err=task_init(child_task,  "fork program", 0, frame->eip,
                        frame->esp + sizeof(uint32_t)*SYSCALL_PARAM_COUNT);
    
    if(err < 0){
        task_free(child_task);
        log_printf("init task fails\n");
        return -1;
    }

    tss_t * tss = &child_task->tss;
    tss->eax = 0;                       // 子进程返回0
    tss->ebx = frame->ebx;
    tss->ecx = frame->ecx;
    tss->edx = frame->edx;
    tss->esi = frame->esi;
    tss->edi = frame->edi;
    tss->ebp = frame->ebp;

    tss->cs = frame->cs;
    tss->ds = frame->ds;
    tss->es = frame->es;
    tss->fs = frame->fs;
    tss->gs = frame->gs;
    tss->eflags = frame->eflags;
    child_task->parent_task = parent_task;
    child_task->tss.cr3 =parent_task->tss.cr3;

    // 复制父进程的内存空间到子进程
    if ((child_task->tss.cr3 = memory_copy_uvm(parent_task->tss.cr3)) <= 0) {
        task_free(child_task);
        log_printf("init task fails\n");
        return -1;
    }
    //将进程加入到就绪队列中
    task_start(child_task);
    return child_task->pid;
};
static int load_phdr(int file, Elf32_Phdr * phdr, uint32_t page_dir) {
    //将elf的一块加载到内存里

    // 生成的ELF文件要求是页边界对齐的
    ASSERT((phdr->p_vaddr & (MEM_PAGE_SIZE - 1)) == 0);

    // 分配空间
    uint32_t err = memory_alloc_for_page_dir(page_dir, phdr->p_vaddr, phdr->p_memsz, PTE_P | PTE_U | PTE_W);
    if (err == 0) {
        log_printf("no memory\n");
        return -1;
    }

    // 调整当前的读写位置
    if (sys_lseek(file, phdr->p_offset, 0) < 0) {
        log_printf("read file failed\n");
        return -1;
    }

    uint32_t vaddr = phdr->p_vaddr;
    uint32_t size = phdr->p_filesz;
    while (size > 0) {
        //获取page_dir的物理地址，因为还没有切换cr3，所以只能用物理地址
        //要分块读取，因为物理地址不一定连续
        uint32_t phy_addr=memory_get_paddr(page_dir,vaddr);
        int read_size=(size > MEM_PAGE_SIZE)?(MEM_PAGE_SIZE):size;
        if(sys_read(file,(void *)phy_addr,read_size)<0){
            log_printf("read file failed\n");
            return -1;
        }
        size-=read_size;
        vaddr+=read_size;
    }
    return 0;
}

static uint32_t load_elf_file (task_t * task, const char * name, uint32_t page_dir){
    //将elf文件加载到内存中,返回入口地址
    Elf32_Ehdr elf_hdr;
    Elf32_Phdr elf_phdr;

    // 以只读方式打开
    int file = sys_open(name, 0);   // todo: flags暂时用0替代
    if (file < 0) {
        log_printf("open file failed.%s\n", name);
        goto load_failed;
    }

    // 先读取文件头
    int cnt = sys_read(file, (char *)&elf_hdr, sizeof(Elf32_Ehdr));
    if (cnt < sizeof(Elf32_Ehdr)) {
        log_printf("elf hdr too small. size=%d\n", cnt);
        goto load_failed;
    }

    // 做点必要性的检查。当然可以再做其它检查
    if ((elf_hdr.e_ident[0] != ELF_MAGIC) || (elf_hdr.e_ident[1] != 'E')
        || (elf_hdr.e_ident[2] != 'L') || (elf_hdr.e_ident[3] != 'F')) {
        log_printf("check elf indent failed.\n");
        goto load_failed;
    }

    // 必须是可执行文件和针对386处理器的类型，且有入口
    if ((elf_hdr.e_type != ET_EXEC) || (elf_hdr.e_machine != ET_386) || (elf_hdr.e_entry == 0)) {
        log_printf("check elf type or entry failed.\n");
        goto load_failed;
    }

    // 必须有程序头部
    if ((elf_hdr.e_phentsize == 0) || (elf_hdr.e_phoff == 0)) {
        log_printf("none programe header\n");
        goto load_failed;
    }

    // 然后从中加载程序头，将内容拷贝到相应的位置
    uint32_t e_phoff = elf_hdr.e_phoff;
    for (int i = 0; i < elf_hdr.e_phnum; i++, e_phoff += elf_hdr.e_phentsize) {
        if (sys_lseek(file, e_phoff, 0) < 0) {
            log_printf("read file failed\n");
            goto load_failed;
        }

        // 读取程序头后解析，这里不用读取到新进程的页表中，因为只是临时使用下
        cnt = sys_read(file, (char *)&elf_phdr, sizeof(Elf32_Phdr));
        if (cnt < sizeof(Elf32_Phdr)) {
            log_printf("read file failed\n");
            goto load_failed;
        }

        // 简单做一些检查，如有必要，可自行加更多
        // 主要判断是否是可加载的类型，并且要求加载的地址必须是用户空间
        if ((elf_phdr.p_type != PT_LOAD) || (elf_phdr.p_vaddr < MEMORY_TASK_BASE)) {
           continue;
        }

        // 加载当前程序头
        int err = load_phdr(file, &elf_phdr, page_dir);
        if (err < 0) {
            log_printf("load program hdr failed\n");
            goto load_failed;
        }

        // 简单起见，不检查了，以最后的地址为bss的地址
        task->heap_start = elf_phdr.p_vaddr + elf_phdr.p_memsz;
        task->heap_end = task->heap_start;

   }

    sys_close(file);
    return elf_hdr.e_entry;

load_failed:
    if (file >= 0) {
        sys_close(file);
    }

    return 0;
}

static int copy_args (char * to, uint32_t page_dir, int argc, char **argv) {
    // 在stack_top中依次写入argc, argv指针，参数字符串
    task_args_t task_args;
    task_args.argc = argc;
    task_args.argv = (char **)(to + sizeof(task_args_t));

    // 复制各项参数, 跳过task_args和参数表
    // 各argv参数写入的内存空间
    char * dest_arg = to + sizeof(task_args_t) + sizeof(char *) * (argc);   // 留出结束符
    
    // argv表
    char ** dest_argv_tb = (char **)memory_get_paddr(page_dir, (uint32_t)(to + sizeof(task_args_t)));
    ASSERT(dest_argv_tb != 0);

    for (int i = 0; i < argc; i++) {
        char * from = argv[i];

        // 不能用kernel_strcpy，因为to和argv不在一个页表里
        int len = kernel_strlen(from) + 1;   // 包含结束符
        int err = memory_copy_uvm_data((uint32_t)dest_arg, page_dir, (uint32_t)from, len);
        ASSERT(err >= 0);

        // 关联ar
        dest_argv_tb[i] = dest_arg;

        // 记录下位置后，复制的位置前移
        dest_arg += len;
    }

     // 写入task_args
    return memory_copy_uvm_data((uint32_t)to, page_dir, (uint32_t)&task_args, sizeof(task_args_t));
}

int sys_execve(const char *name, char * const *argv, char * const *env){

    task_t *current_task=get_task_current();
    kernel_memcpy(current_task->name,get_file_name(name),TASK_NAME_SIZE);

    //创建新的页表
    uint32_t old_page_dir=current_task->tss.cr3;
    uint32_t new_page_dir=memory_create_uvm();
    if(new_page_dir==0){
        goto execve_fail;
    }
    //加载elf文件
    uint32_t entry=load_elf_file(current_task,name,new_page_dir);
    if(entry==0){
        goto execve_fail;
    }
    //为进程创建新的栈空间
    uint32_t stack_top = MEM_TASK_STACK_TOP - MEM_TASK_ARG_SIZE;    // 预留一部分参数空间
    
    uint32_t err=memory_alloc_for_page_dir(new_page_dir,MEM_TASK_STACK_TOP - MEM_TASK_STACK_SIZE,MEM_TASK_STACK_SIZE,PTE_P | PTE_U | PTE_W);
    if(err==0){
        goto execve_fail;
    }

    // 复制参数，写入到栈顶的后边
    int argc = strings_count(argv);
    int err1 = copy_args((char *)stack_top, new_page_dir, argc, argv);
    if (err1 < 0) {
        goto execve_fail;
    }

    //修改新进程的寄存器
    syscall_frame_t * frame = (syscall_frame_t *)(current_task->tss.esp0 - sizeof(syscall_frame_t));
    frame->eip = entry;
    frame->eax = frame->ebx = frame->ecx = frame->edx = 0;
    frame->esi = frame->edi = frame->ebp = 0;
    frame->eflags = EFLAGS_DEFAULT| EFLAGS_IF;  // 段寄存器无需修改


    // 内核栈不用设置，保持不变，后面调用memory_destroy_uvm并不会销毁内核栈的映射。
    // 但用户栈需要更改, 同样要加上调用门的参数压栈空间
    frame->esp = stack_top - sizeof(uint32_t)*SYSCALL_PARAM_COUNT;

    //设置新的cr3 页表
    current_task->tss.cr3=new_page_dir;
    mmu_set_page_dir(new_page_dir);

    //释放旧页表
    memory_destroy_uvm(old_page_dir);
    return 0;
execve_fail:
    if(new_page_dir){
        memory_destroy_uvm(new_page_dir);
        current_task->tss.cr3=old_page_dir;
        mmu_set_page_dir(old_page_dir);
    }
    return -1;
};

//根据task的文件指针，返回其对应的file_t结构
file_t * task_file (int fd){
    if( fd < 0 || fd >=TASK_OFILE_NR){
        return (file_t *)0;
    }
    return get_task_current()->file_table[fd];
};
//在task_t的表格里，分配一个空闲项，返回他的句柄
int task_alloc_fd (file_t * file){
    task_t *current_task=get_task_current();
    for(int i=0;i<TASK_OFILE_NR;i++){
        if(current_task->file_table[i]==(file_t *)0){
            current_task->file_table[i]=file;
            return i;
        }
    }
    return -1;
};
//在task_t的表格里，删除一个空闲项
void task_remove_fd (int fd){
    if( fd < 0 || fd >=TASK_OFILE_NR){
        return ;
    }
    get_task_current()->file_table[fd]=(file_t *)0;
};